import argparse
import asyncio
from pathlib import Path
import yaml
import logging

from utils.logger import LoggerManager
from master.master_main import MasterNode
from worker.worker_main import Worker

def load_config(config_path: str) -> dict:
    """Load configuration from YAML file.
    
    Args:
        config_path: Path to configuration file
        
    Returns:
        dict: Configuration parameters
    """
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def setup_environment(config: dict) -> None:
    """Setup runtime environment.
    
    Args:
        config: Configuration parameters
    """
    # Create necessary directories
    Path(config['logging']['log_dir']).mkdir(parents=True, exist_ok=True)
    if 'cache_dir' in config:
        Path(config['cache_dir']).mkdir(parents=True, exist_ok=True)

async def run_master(config_path: str) -> None:
    """Run master node.
    
    Args:
        config_path: Path to configuration file
    """
    master = MasterNode(config_path)
    await master.start()

async def run_worker(config_path: str) -> None:
    """Run worker node.
    
    Args:
        config_path: Path to configuration file
    """
    worker = Worker(config_path)
    await worker.start()

def main():
    """Main entry point for the distributed system."""
    parser = argparse.ArgumentParser(
        description='Start DeepSlicing distributed system'
    )
    parser.add_argument('--config', type=str, required=True,
                      help='Path to configuration file')
    parser.add_argument('--role', type=str, required=True,
                      choices=['master', 'worker'],
                      help='Node role (master or worker)')
    
    args = parser.parse_args()
    
    # Load configuration
    config = load_config(args.config)
    
    # Setup environment
    setup_environment(config)
    
    # Initialize logging
    logger_manager = LoggerManager(config['logging'])
    logger = logger_manager.get_logger('DeepSlicing')
    logger.info(f"Starting {args.role} node...")
    
    try:
        # Run appropriate node type
        if args.role == 'master':
            asyncio.run(run_master(args.config))
        else:  # worker
            asyncio.run(run_worker(args.config))
    except KeyboardInterrupt:
        logger.info("Shutting down...")
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        raise

if __name__ == '__main__':
    main()